#include <iostream>
#include <stdlib.h>
#include <cmath>
#include <random>
#include <numeric>

#include "gtest/gtest.h"

#include "kpextVkContext.h"
#include "KpExtGP.h"
#include "kpextOpImgF32C1Average.h"

#include "kpextCtxGlobal.h"

using namespace std;

TEST(kpextFuncTest, OpImgF32C1Average){
    ASSERT_NE(nullptr, GvkContextPtr);

    int cols = 1024, rows = 1000;

    auto cmdbuf = GvkContextPtr->cmdbuf();

    size_t size = static_cast<size_t>(cols * rows);
    auto stagebuf1 = GvkContextPtr->stagingBuf(size * 4);
    auto stagebuf2 = GvkContextPtr->stagingBuf(size * 4);
    auto inbuf = GvkContextPtr->deviceBuf(size * 4);
    auto outbuf = GvkContextPtr->deviceBuf(size * 4);
    auto op = std::make_shared<kpext::OpImgF32C1Average>(
                GvkContextPtr
                );

    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, outbuf);
    ASSERT_NE(nullptr, op);

    {
        std::random_device r;

        std::default_random_engine e1(r());
        std::normal_distribution<float> normal_dist(
                    0,  // mean
                    10  // sigma
                    );

        auto data = stagebuf1->data<float>();
        for (size_t i = 0; i < size; ++i) {
            data[i] = normal_dist(e1);
        }
        stagebuf1->unmap();
    }


    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;
    kpext::CmdMemoryBarrier barrier;


    barrierTx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );
    barrier.prepare(
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    op->rebuild(cols, rows,
                1.0f / 3.0f,
                inbuf, outbuf
                );

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1, size * 4);
    cmdbuf->record(&bufcpy);
    bufset.prepare(outbuf, 0);
    cmdbuf->record(&bufset);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());
    cmdbuf->record(&barrier);

    uint32_t v(0);
    BITCAST_ASSIGN(v, 2.973f);
    bufset.prepare(inbuf, v);
    cmdbuf->record(&bufset);
    cmdbuf->record(&barrier);

    cmdbuf->record(op.get());
    cmdbuf->record(&barrier);

    BITCAST_ASSIGN(v, 3.531f);
    bufset.prepare(inbuf, v);
    cmdbuf->record(&bufset);
    cmdbuf->record(&barrier);

    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->submit();
    cmdbuf->waitExec();

    {
        auto p1 = stagebuf1->data<float>();
        auto p2 = stagebuf2->data<float>();
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf1->vector<float>());
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<float>());
        float maxerr = 0;
        for (int r = 0; r < rows; ++r) {
            for (int c = 0; c < cols; ++c) {
                int idx = c + r * cols;
                float f1 = p1[idx];

                f1 = f1 + 2.973f;
                f1 = f1 + 3.531f;
                f1 = f1 / 3.0f;

                float f2 = p2[idx];

                f2 = std::fabsf(f1 - f2);
                if(f2 > maxerr){
                    maxerr = f2;
                }
            }
//            auto sptr = p2 + r * cols;
//            std::vector<float> rowdisp(sptr, sptr + cols);
//            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf1->unmap();
        stagebuf2->unmap();
        std::cout << "max error = " << maxerr << std::endl;
        EXPECT_LE(maxerr, 1e-5f);
    }

}


GTEST_API_ int main(int argc, char **argv) {
    printf("Running main() from %s\n", __FILE__);
    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_DEBUG);
    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    //  kpextSetLoggerPattern("[%D %H:%M:%S.%e] [%^%l%$] %v");
    kpextSetLoggerPattern("   [SPDLOG_%l] %v");

    testing::InitGoogleTest(&argc, argv);

    testing::Environment* const foo_env =
            testing::AddGlobalTestEnvironment(new KpExtCtxEnv);
    UNUSE_VAR(foo_env);

    return RUN_ALL_TESTS();
}
